package com.cyzy.caster;

import org.springframework.web.socket.WebSocketSession;

import javax.websocket.Session;
import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 弹幕消息广播器，支持按班级ID分组广播
 */
public class DanmakuBroadcaster {
    // 存储班级ID与原生WebSocket连接的映射 (班级ID -> 连接集合)
    private static final Map<String, Set<Session>> classSessions =
            Collections.synchronizedMap(new HashMap<>());

    // 存储班级ID与Spring WebSocket连接的映射 (班级ID -> 连接集合)
    private static final Map<String, Set<WebSocketSession>> classSpringSessions =
            Collections.synchronizedMap(new HashMap<>());

    // 存储连接与班级ID的反向映射 (用于断开连接时清理)
    private static final Map<Session, String> sessionToClass = new ConcurrentHashMap<>();
    private static final Map<WebSocketSession, String> springSessionToClass = new ConcurrentHashMap<>();

    // 用于去重的消息ID缓存
    private static final ConcurrentHashMap<String, Boolean> broadcastedMessages = new ConcurrentHashMap<>();


    /**
     * 添加原生WebSocket连接到指定班级
     */
    public static void addSession(Session session, String classId) {
        if (session != null && session.isOpen() && classId != null) {
            // 按班级ID分组存储连接
            classSessions.computeIfAbsent(classId, k ->
                    Collections.synchronizedSet(new HashSet<>())
            ).add(session);
            // 记录连接所属班级
            sessionToClass.put(session, classId);
        }
    }

    /**
     * 添加Spring WebSocket连接到指定班级
     */
    public static void addSpringSession(WebSocketSession session, String classId) {
        if (session != null && session.isOpen() && classId != null) {
            classSpringSessions.computeIfAbsent(classId, k ->
                    Collections.synchronizedSet(new HashSet<>())
            ).add(session);
            springSessionToClass.put(session, classId);
        }
    }

    /**
     * 移除原生WebSocket连接
     */
    public static void removeSession(Session session) {
        if (session != null) {
            String classId = sessionToClass.remove(session);
            if (classId != null) {
                Set<Session> sessions = classSessions.get(classId);
                if (sessions != null) {
                    sessions.remove(session);
                    // 清理空班级集合
                    if (sessions.isEmpty()) {
                        classSessions.remove(classId);
                    }
                }
            }
        }
    }

    /**
     * 移除Spring WebSocket连接
     */
    public static void removeSpringSession(WebSocketSession session) {
        if (session != null) {
            String classId = springSessionToClass.remove(session);
            if (classId != null) {
                Set<WebSocketSession> sessions = classSpringSessions.get(classId);
                if (sessions != null) {
                    sessions.remove(session);
                    if (sessions.isEmpty()) {
                        classSpringSessions.remove(classId);
                    }
                }
            }
        }
    }

    /**
     * 向指定班级的所有客户端广播消息
     * @param classId 班级ID (null则广播给所有班级)
     * @param message 消息内容
     */
    public static void broadcast(String classId, String message) {
        if (message == null || message.trim().isEmpty()) {
            return;
        }

        // 生成消息ID去重
        String messageId = message.hashCode() + "_" + System.currentTimeMillis();
        if (broadcastedMessages.putIfAbsent(messageId, Boolean.TRUE) != null) {
            System.out.println("重复消息跳过: " + messageId);
            return;
        }

        try {
            // 广播原生WebSocket连接
            if (classId == null) {
                // 广播所有班级
                classSessions.values().forEach(sessions -> broadcastToSessions(sessions, message));
            } else {
                // 只广播指定班级
                Set<Session> sessions = classSessions.get(classId);
                if (sessions != null) {
                    broadcastToSessions(sessions, message);
                }
            }

            // 广播Spring WebSocket连接
            if (classId == null) {
                classSpringSessions.values().forEach(sessions -> broadcastToSpringSessions(sessions, message));
            } else {
                Set<WebSocketSession> sessions = classSpringSessions.get(classId);
                if (sessions != null) {
                    broadcastToSpringSessions(sessions, message);
                }
            }
        } finally {
            broadcastedMessages.remove(messageId);
        }
    }

    // 向原生WebSocket会话集合广播消息
    private static void broadcastToSessions(Set<Session> sessions, String message) {
        sessions.forEach(session -> {
            if (session != null && session.isOpen()) {
                try {
                    session.getBasicRemote().sendText(message);
                } catch (IOException e) {
                    System.err.println("广播失败: " + e.getMessage());
                    try { session.close(); } catch (IOException ex) {}
                }
            }
        });
    }

    // 向Spring WebSocket会话集合广播消息
    private static void broadcastToSpringSessions(Set<WebSocketSession> sessions, String message) {
        sessions.forEach(session -> {
            if (session != null && session.isOpen()) {
                try {
                    session.sendMessage(new org.springframework.web.socket.TextMessage(message));
                } catch (Exception e) {
                    System.err.println("广播失败: " + e.getMessage());
                    try { session.close(); } catch (IOException ex) {}
                }
            }
        });
    }
}